/*
*	目的：这段代码实现了一个基于 令牌桶算法 的速率限制器，常用于流量控制、带宽管理等场景。
*		 令牌桶算法通过控制  令牌桶的生成速率 和 桶的容量 来限制某些操作的执行速率。	
*
*该算法可以让系统允许短时间内超出速率限制的请求，但会根据令牌的可用情况进行调控。
*/

#include "tbf.h"
#include <signal.h>
#include <stdlib.h>
#include <unistd.h>
#include <sys/time.h>
#include <stdio.h>

//定义结构体
typedef struct {
	int cps;//速率，每秒生成多少令牌
	int token;//当前剩余的令牌数
	int burst;//令牌桶的容量上限（超出部分将会被丢弃）
}tbf_t;

//全局变量
static tbf_t *tbflibs[TBFMAX];    //存储所有令牌同实例 
static int inited = 0;            //标记是否已经初始化（初始化只需要做一次）

static int _tbf_value(void);
static void _alarm_handler(int s);
static int __get_free_pos(void);

//第一步： 初始化
//初始化一个令牌桶实例，设置速率和容量，并为其分配内存空间
//参数1：每秒生成的令牌数
//参数2：令牌桶的容量
int tbf_init(int cps, int burst)
{
	int pos;

	//如果令牌桶系统还没有初始化，就调用_tbf_value()函数初始化定时器，开始定时生成令牌
	if(!inited) {
		_tbf_value();
		inited = 1;
	}
	//动态分配内存为新的令牌桶（tbf_t）,并设置其 cps 和 burst.
	tbf_t *mytbf = malloc(sizeof(tbf_t));
	if(NULL == mytbf) {
		//判断是否开辟成功
		return -1;
	}
	mytbf->cps = cps;
	mytbf->token = 0;
	mytbf->burst = burst;

	//通过__get_free_pos()获取一个空位，将新创建的令牌桶存入 tbflibs 数组中
	pos = __get_free_pos();
	if(-1 == pos) {
		free(mytbf);
		return -1;
	}
	tbflibs[pos] = mytbf;
	//成功返回令牌桶的索引，失败返回-1.
	return pos;
}


//查找 tbflibs 数组中的空位置，用于存储新的令牌桶
static int __get_free_pos(void)
{
	int i;
	//遍历 tbflibs 数组
	for(i = 0; i < TBFMAX; i++) {
		//如果找到一个 NULL 位置，返回该位置的索引
		if(tbflibs[i] == NULL) {
			return i;
		}
	}
	//成功返回位置的索引，失败返回-1，表示没有可用空间
	return -1;
}


//第二步： 定时生成令牌
//使用定时器每秒生成令牌，并将其加到每个令牌桶中
//初始化定时器，定时调用_alarm_handler函数来生成令牌
static int _tbf_value(void)
{
	struct sigaction act;
	struct itimerval itv;

	act.sa_handler = _alarm_handler;
	act.sa_flags = 0;
	sigemptyset(&act.sa_mask);
	sigaction(SIGALRM, &act, NULL);

	itv.it_interval.tv_sec = 1;
	itv.it_interval.tv_usec = 0;
	itv.it_value.tv_sec = 1;
	itv.it_value.tv_usec = 0;
	//使用setitimer设置一个定时器，每秒触发依次 SIGALRM 信号
	setitimer(ITIMER_REAL, &itv, NULL);
	
	return 0;
}


//处理 SIGALARM 信号，每秒更新所有令牌桶的令牌数量（每秒钟会被定时器调用一次）
static void _alarm_handler(int s)
{
	int i;
	//遍历 tbflibs 数组中的每个令牌桶
	for(i = 0; i < TBFMAX; i++) {
		if(tbflibs[i]) {
			//增加 cps 数量的令牌
			tbflibs[i]->token += tbflibs[i]->cps;
			//如果令牌数量(token) 超过 桶的容量(burst),则将其限制为 桶的容量(burst)
			if(tbflibs[i]->token > tbflibs[i]->burst)
				tbflibs[i]->token = tbflibs[i]->burst;
		}
	}
}

//第三步： 获取令牌
//从指定的令牌桶中 获取令牌，令牌的数量由桶的速率控制，超出时挂起等待。（核心部分：用于控制操作的速率）
//参数1：令牌桶的索引 td
//参数2：需要的令牌数量 ntoken
int tbf_fetch_token(int td, int ntoken)
{
	// 如果当前令牌桶的令牌数为 0，就进入while循环，调用pause()函数挂起当前进程，直到令牌桶中有可用令牌
	while(tbflibs[td]->token == 0) {
		pause();
	}

	//如果请求的令牌数 大于 当前的可用令牌
	if(ntoken > tbflibs[td]->token) {
		//就取所有可用的令牌
		ntoken = tbflibs[td]->token;
	}
	//从令牌桶中 减去 获取的令牌数
	tbflibs[td]->token -= ntoken;
	//返回实际获取的令牌数
	return ntoken;
}


//第四步： 销毁
//销毁指定的令牌桶实例，释放相关资源
void tbf_destroy(int td)
{
	//释放指定令牌桶实例的内存
	free(tbflibs[td]);
	//将 tbflibs 数组中对应位置的指针置为 NULL
	tbflibs[td] = NULL;
}
